from runner import Runner
from common.arguments import get_common_args, get_coma_args, get_mixer_args
import gym

if __name__ == '__main__':
    for i in range(8):
        args = get_common_args()
        if args.alg.find('coma') > -1:
            args = get_coma_args(args)
        else:
            args = get_mixer_args(args)
        env_name = "intersection-multi-agent-v0"
        # env_name = "highway-multi-agent-v0"
        env = gym.make(env_name)
        args.map = env_name
        env_info = env.get_env_info()
        # print(env_info)
        args.n_actions = env_info["n_actions"]
        args.n_agents = env_info["n_agents"]
        args.state_shape = env_info["state_shape"]
        args.obs_shape = env_info["obs_shape"]
        args.episode_limit = env_info["episode_limit"]
        runner = Runner(env, args)
        if not args.evaluate:
            runner.run(i)
        else:
            episode_rewards = runner.evaluate()
            print('The episode_rewards of {} is {}'.format(args.alg, episode_rewards))
            break
        env.close()

